﻿Imports System.ComponentModel

Public Enum SelectionType
    清空
    反选
    全选
End Enum

<ProvideProperty("SelectionSource", GetType(Control))>
<ProvideProperty("SelectionType", GetType(Control))>
Public Class FastSelect
    Inherits Component
    Implements IExtenderProvider

    Private sourceControlsDict As Dictionary(Of Control, Control)
    Private typeDict As Dictionary(Of Control, SelectionType)

    Public Sub New()
        Me.New(Nothing)
    End Sub

    Public Sub New(ByVal container As IContainer)
        If container IsNot Nothing Then container.Add(Me)
        sourceControlsDict = New Dictionary(Of Control, Control)
        typeDict = New Dictionary(Of Control, SelectionType)
    End Sub

    Public Function CanExtend(ByVal extendee As Object) As Boolean Implements System.ComponentModel.IExtenderProvider.CanExtend
        If extendee Is Nothing Then Return False
        If TypeOf extendee Is Button OrElse TypeOf extendee Is Label OrElse TypeOf extendee Is PictureBox Then Return True
        Return False
    End Function

    <Category("速选"), Description("速选源控件"), Localizable(True)>
    Public Function GetSelectionSource(ByVal control As Control) As Control
        If sourceControlsDict.ContainsKey(control) Then Return sourceControlsDict(control) Else Return Nothing
    End Function

    Public Sub SetSelectionSource(ByVal control As Control, ByVal selectionSource As Control)
        If selectionSource IsNot Nothing Then
            If sourceControlsDict.ContainsKey(control) AndAlso sourceControlsDict(control) IsNot selectionSource Then
                SetSelectionSource(control, Nothing)
            End If
            sourceControlsDict.Add(control, selectionSource)
            typeDict.Add(control, SelectionType.清空)
            AddHandler control.Click, AddressOf control_Click
        Else
            sourceControlsDict.Remove(control)
            typeDict.Remove(control)
            RemoveHandler control.Click, AddressOf control_Click
        End If
    End Sub

    <Category("速选"), Description("速选方式"), DefaultValue(SelectionType.清空), Localizable(True)>
    Public Function GetSelectionType(ByVal control As Control) As SelectionType
        If typeDict.ContainsKey(control) Then
            Return typeDict(control)
        Else
            Return SelectionType.清空
        End If
    End Function

    Public Sub SetSelectionType(ByVal control As Control, ByVal selectionType As SelectionType)
        If typeDict.ContainsKey(control) Then
            typeDict(control) = selectionType
        Else
            typeDict.Add(control, selectionType)
        End If
    End Sub

    Private Sub control_Click(ByVal sender As Object, ByVal e As EventArgs)
        Dim control As Control = TryCast(sender, Control)
        Dim selectionSource As Control = sourceControlsDict(control)
        Dim selectType As SelectionType = typeDict(control)

        If TypeOf selectionSource Is DataGridView Then
            Dim dataGridView As DataGridView = TryCast(selectionSource, DataGridView)
            For Each row As DataGridViewRow In dataGridView.Rows
                row.Selected = ChangeSelected(row.Selected, selectType)
            Next
        ElseIf TypeOf selectionSource Is CheckedListBox Then
            Dim checkedListBox As CheckedListBox = TryCast(selectionSource, CheckedListBox)
            For i As Integer = 0 To checkedListBox.Items.Count - 1
                Dim _selected As Boolean = checkedListBox.GetItemChecked(i)
                _selected = ChangeSelected(_selected, selectType)
                checkedListBox.SetItemChecked(i, _selected)
                checkedListBox.SetSelected(i, _selected)
            Next
        ElseIf TypeOf selectionSource Is ListBox Then
            Dim listBox As ListBox = TryCast(selectionSource, ListBox)
            For i As Integer = 0 To listBox.Items.Count - 1
                Dim _selected As Boolean = listBox.GetSelected(i)
                _selected = ChangeSelected(_selected, selectType)
                listBox.SetSelected(i, _selected)
            Next
        ElseIf TypeOf selectionSource Is TreeView Then
            Dim treeView As TreeView = TryCast(selectionSource, TreeView)
            For Each treeNode As TreeNode In GetTreeNodes(treeView)
                treeNode.Checked = ChangeSelected(treeNode.Checked, selectType)
            Next
        ElseIf selectionSource.Controls.Count > 0 Then
            For Each subControl As Control In selectionSource.Controls
                If TypeOf subControl Is CheckBox Then
                    Dim checkBox As CheckBox = TryCast(subControl, CheckBox)
                    checkBox.Checked = ChangeSelected(checkBox.Checked, selectType)
                End If
            Next
        End If
    End Sub

    Private Function ChangeSelected(ByVal isSelected As Boolean, ByVal type As SelectionType) As Boolean
        If type = SelectionType.清空 Then
            Return False
        ElseIf type = SelectionType.全选 Then
            Return True
        ElseIf type = SelectionType.反选 Then
            Return Not isSelected
        Else
            Throw New NotImplementedException()
        End If
    End Function

    Private Function GetTreeNodes(ByVal treeView As TreeView) As IEnumerable(Of TreeNode)
        Dim result As New List(Of TreeNode)
        For Each node As TreeNode In treeView.Nodes
            result.Add(node)
            For Each subNode In GetTreeNodes(node)
                result.Add(subNode)
            Next
        Next
        Return result
    End Function

    Private Function GetTreeNodes(ByVal treeNode As TreeNode) As IEnumerable(Of TreeNode)
        Dim result As New List(Of TreeNode)
        For Each node As TreeNode In treeNode.Nodes
            result.Add(node)
            For Each subNode In GetTreeNodes(node)
                result.Add(subNode)
            Next
        Next
        Return result
    End Function
End Class
